#include "Graphics/gpuprefabs.h"


namespace Jin {

	GPUGraphicsPipelinePool::GPUGraphicsPipelinePool(GPUDevice& device) : device(device) {
		states.reserve(GPL_NUM);
		pipelines.reserve(GPL_NUM);
		if (!_CreateShaders()) {
			return;
		}
		_CreateRasterizerStates();
		_CreateMultisampleStates();
		_CreateDepthStencilStates();
		_CreateTargetInfo();
		_CreateColorTargetDescriptions();
		if (!_CreateGraphicsPipelines()) {
			return;
		}
	}

	Uint32 GPUGraphicsPipelinePool::GetPipelineID(GPLState state) {
		if (m_stateCache != state) {
			auto iter = std::find(states.begin(), states.end(), state);
			if (iter != states.end()) {
				m_stateCache = state;
				m_idCache = (Uint32)(iter - states.begin());
				return m_idCache;
			}
			if (!_CreateGraphicsPipelineState(state)) {
				return INVALID_ID;
			}
			m_stateCache = state;
			m_idCache = (Uint32)(states.size() - 1);
		}
		return m_idCache;
	}

	bool GPUGraphicsPipelinePool::_CreateShaders() {
		static const char* VSFilename[] = {
			"PositionVertex.vert",
			"PointVertex.vert",
			"LineVertex.vert",
			"CircleVertex.vert",
			"TextureVertex.vert"
		};
		static const char* FSFilename[] = {
			"OutColor.frag",
			"CirclePoint.frag",
			"AACirclePoint.frag",
			"Line.frag",
			"AALine.frag",
			"Circle.frag",
			"AACircle.frag",
			"OutTexture.frag",
			"DotTextureColor.frag"
		};

		SDLFileStorage storage("shaders");
		if (!storage) {
			return false;
		}
		string ext;
		stringv driver = device.GetDeviceDriver();
		stringv entrypoint = "main";
		GPUShaderFormat format = SDL_GPU_SHADERFORMAT_INVALID;
		if (driver == "vulkan") {
			format = SDL_GPU_SHADERFORMAT_SPIRV;
			ext = ".spv";
		}
		else if (driver == "direct3d12") {
			format = SDL_GPU_SHADERFORMAT_DXIL;
			ext = ".dxil";
		}
		else {
			entrypoint = "main0";
			format = SDL_GPU_SHADERFORMAT_METALLIB;
			ext = ".metallib";
		}

		auto& vertexShaders = statePrefabs.vertexShaders;
		auto& fragmentShaders = statePrefabs.fragmentShaders;
		vertexShaders.reserve(VS_NUM);
		fragmentShaders.reserve(FS_NUM);
		GPUVertexShaderInfo vertexShaderInfo;
		vertexShaderInfo.entrypoint = entrypoint;
		vertexShaderInfo.format = format;
		GPUFragmentShaderInfo fragmentShaderInfo;
		fragmentShaderInfo.entrypoint = entrypoint;
		fragmentShaderInfo.format = format;

		// VS_POSITIONVERTEX
		{
			vertexShaderInfo.numUniformBuffers = 1;		// push MVPVertexUniform
			if (!storage.ReadFile(VSFilename[VS_POSITIONVERTEX] + ext, vertexShaderInfo.code)) {
				return false;
			}
			auto& vertexShader = vertexShaders.emplace_back(std::make_unique<GPUVertexShader>(device, vertexShaderInfo));
			if (!*vertexShader) {
				vertexShaders.pop_back();
				return false;
			}
			vertexShader->vertexBufferDescriptions.resize(1);
			vertexShader->vertexBufferDescriptions[0].pitch = sizeof(PositionVertex);
			vertexShader->vertexAttributes.resize(2);
			vertexShader->vertexAttributes[0].location = 0;
			vertexShader->vertexAttributes[0].format = SDL_GPU_VERTEXELEMENTFORMAT_FLOAT3;
			vertexShader->vertexAttributes[0].offset = offsetof(PositionVertex, position);
			vertexShader->vertexAttributes[1].location = 1;
			vertexShader->vertexAttributes[1].format = SDL_GPU_VERTEXELEMENTFORMAT_FLOAT4;
			vertexShader->vertexAttributes[1].offset = offsetof(PositionVertex, color);
		}
		// VS_POINTVERTEX
		{
			vertexShaderInfo.numUniformBuffers = 2;		// push MVPVertexUniform, PointVertexUniform
			if (!storage.ReadFile(VSFilename[VS_POINTVERTEX] + ext, vertexShaderInfo.code)) {
				return false;
			}
			auto& vertexShader = vertexShaders.emplace_back(std::make_unique<GPUVertexShader>(device, vertexShaderInfo));
			if (!*vertexShader) {
				vertexShaders.pop_back();
				return false;
			}
			vertexShader->vertexBufferDescriptions.resize(1);
			vertexShader->vertexBufferDescriptions[0].pitch = sizeof(PositionVertex);
			vertexShader->vertexAttributes.resize(2);
			vertexShader->vertexAttributes[0].location = 0;
			vertexShader->vertexAttributes[0].format = SDL_GPU_VERTEXELEMENTFORMAT_FLOAT3;
			vertexShader->vertexAttributes[0].offset = offsetof(PositionVertex, position);
			vertexShader->vertexAttributes[1].location = 1;
			vertexShader->vertexAttributes[1].format = SDL_GPU_VERTEXELEMENTFORMAT_FLOAT4;
			vertexShader->vertexAttributes[1].offset = offsetof(PositionVertex, color);
		}
		// VS_LINEVERTEX
		{
			vertexShaderInfo.numUniformBuffers = 2;		// push MVPVertexUniform, LineVertexUniform
			vertexShaderInfo.numStorageBuffers = 1;		// bind LineVertex
			if (!storage.ReadFile(VSFilename[VS_LINEVERTEX] + ext, vertexShaderInfo.code)) {
				return false;
			}
			auto& vertexShader = vertexShaders.emplace_back(std::make_unique<GPUVertexShader>(device, vertexShaderInfo));
			if (!*vertexShader) {
				vertexShaders.pop_back();
				return false;
			}
		}
		// VS_CIRCLEVERTEX
		{
			vertexShaderInfo.numUniformBuffers = 2;		// push MVPVertexUniform, CircleVertexUniform
			vertexShaderInfo.numStorageBuffers = 1;		// bind CircleVertex
			if (!storage.ReadFile(VSFilename[VS_CIRCLEVERTEX] + ext, vertexShaderInfo.code)) {
				return false;
			}
			auto& vertexShader = vertexShaders.emplace_back(std::make_unique<GPUVertexShader>(device, vertexShaderInfo));
			if (!*vertexShader) {
				vertexShaders.pop_back();
				return false;
			}
		}
		// VS_TEXTUREVERTEX
		{
			vertexShaderInfo.numUniformBuffers = 2;		// push MVPVertexUniform, glm::mat4
			vertexShaderInfo.numStorageBuffers = 0;
			if (!storage.ReadFile(VSFilename[VS_TEXTUREVERTEX] + ext, vertexShaderInfo.code)) {
				return false;
			}
			auto& vertexShader = vertexShaders.emplace_back(std::make_unique<GPUVertexShader>(device, vertexShaderInfo));
			if (!*vertexShader) {
				vertexShaders.pop_back();
				return false;
			}
			vertexShader->vertexBufferDescriptions.resize(1);
			vertexShader->vertexBufferDescriptions[0].pitch = sizeof(TextureVertex);
			vertexShader->vertexAttributes.resize(3);
			vertexShader->vertexAttributes[0].location = 0;
			vertexShader->vertexAttributes[0].format = SDL_GPU_VERTEXELEMENTFORMAT_FLOAT3;
			vertexShader->vertexAttributes[0].offset = offsetof(TextureVertex, position);
			vertexShader->vertexAttributes[1].location = 1;
			vertexShader->vertexAttributes[1].format = SDL_GPU_VERTEXELEMENTFORMAT_FLOAT4;
			vertexShader->vertexAttributes[1].offset = offsetof(TextureVertex, color);
			vertexShader->vertexAttributes[2].location = 2;
			vertexShader->vertexAttributes[2].format = SDL_GPU_VERTEXELEMENTFORMAT_FLOAT2;
			vertexShader->vertexAttributes[2].offset = offsetof(TextureVertex, uv);
		}

		// FS_OUTCOLOR
		{
			if (!storage.ReadFile(FSFilename[FS_OUTCOLOR] + ext, fragmentShaderInfo.code)) {
				return false;
			}
			if (!*fragmentShaders.emplace_back(std::make_unique<GPUFragmentShader>(device, fragmentShaderInfo))) {
				fragmentShaders.pop_back();
				return false;
			}
		}
		// FS_CIRCLEPOINT
		{
			if (!storage.ReadFile(FSFilename[FS_CIRCLEPOINT] + ext, fragmentShaderInfo.code)) {
				return false;
			}
			if (!*fragmentShaders.emplace_back(std::make_unique<GPUFragmentShader>(device, fragmentShaderInfo))) {
				fragmentShaders.pop_back();
				return false;
			}
		}
		// FS_AACIRCLEPOINT
		{
			if (!storage.ReadFile(FSFilename[FS_AACIRCLEPOINT] + ext, fragmentShaderInfo.code)) {
				return false;
			}
			if (!*fragmentShaders.emplace_back(std::make_unique<GPUFragmentShader>(device, fragmentShaderInfo))) {
				fragmentShaders.pop_back();
				return false;
			}
		}
		// FS_LINE
		{
			if (!storage.ReadFile(FSFilename[FS_LINE] + ext, fragmentShaderInfo.code)) {
				return false;
			}
			if (!*fragmentShaders.emplace_back(std::make_unique<GPUFragmentShader>(device, fragmentShaderInfo))) {
				fragmentShaders.pop_back();
				return false;
			}
		}
		// FS_AALINE
		{
			if (!storage.ReadFile(FSFilename[FS_AALINE] + ext, fragmentShaderInfo.code)) {
				return false;
			}
			if (!*fragmentShaders.emplace_back(std::make_unique<GPUFragmentShader>(device, fragmentShaderInfo))) {
				fragmentShaders.pop_back();
				return false;
			}
		}
		// FS_CIRCLE
		{
			if (!storage.ReadFile(FSFilename[FS_CIRCLE] + ext, fragmentShaderInfo.code)) {
				return false;
			}
			if (!*fragmentShaders.emplace_back(std::make_unique<GPUFragmentShader>(device, fragmentShaderInfo))) {
				fragmentShaders.pop_back();
				return false;
			}
		}
		// FS_AACIRCLE
		{
			if (!storage.ReadFile(FSFilename[FS_AACIRCLE] + ext, fragmentShaderInfo.code)) {
				return false;
			}
			if (!*fragmentShaders.emplace_back(std::make_unique<GPUFragmentShader>(device, fragmentShaderInfo))) {
				fragmentShaders.pop_back();
				return false;
			}
		}
		// FS_OUTTEXTURE
		{
			fragmentShaderInfo.numSamplers = 1;		// BindFragmentSampler
			if (!storage.ReadFile(FSFilename[FS_OUTTEXTURE] + ext, fragmentShaderInfo.code)) {
				return false;
			}
			if (!*fragmentShaders.emplace_back(std::make_unique<GPUFragmentShader>(device, fragmentShaderInfo))) {
				fragmentShaders.pop_back();
				return false;
			}
		}
		// FS_DOTTEXTURECOLOR
		{
			fragmentShaderInfo.numSamplers = 1;		// BindFragmentSampler
			if (!storage.ReadFile(FSFilename[FS_DOTTEXTURECOLOR] + ext, fragmentShaderInfo.code)) {
				return false;
			}
			if (!*fragmentShaders.emplace_back(std::make_unique<GPUFragmentShader>(device, fragmentShaderInfo))) {
				fragmentShaders.pop_back();
				return false;
			}
		}
		return true;
	}

	void GPUGraphicsPipelinePool::_CreateRasterizerStates() {
		statePrefabs.rasterizerStates.reserve(RS_NUM);
		statePrefabs.rasterizerStates.emplace_back();
	}

	void GPUGraphicsPipelinePool::_CreateMultisampleStates() {
		statePrefabs.multisampleStates.reserve(MS_NUM);
		statePrefabs.multisampleStates.emplace_back();
	}

	void GPUGraphicsPipelinePool::_CreateDepthStencilStates() {
		statePrefabs.depthStencilStates.reserve(DSS_NUM);
		statePrefabs.depthStencilStates.emplace_back();
	}

	void GPUGraphicsPipelinePool::_CreateTargetInfo() {
		statePrefabs.targetInfos.reserve(TI_NUM);
		statePrefabs.targetInfos.emplace_back();
	}

	void GPUGraphicsPipelinePool::_CreateColorTargetDescriptions() {
		statePrefabs.colorTargetDescriptionsPool.reserve(CTD_NUM);
		auto init = std::initializer_list{ GPUColorTargetDescription{.format = device.swapchainTextureFormat } };
		// BLENDMODE_NONE
		{
			auto& desc = statePrefabs.colorTargetDescriptionsPool.emplace_back(init);
			//descs[0].blend_state = {
			//	SDL_GPU_BLENDFACTOR_ONE,
			//	SDL_GPU_BLENDFACTOR_ZERO,
			//	SDL_GPU_BLENDOP_ADD,
			//	SDL_GPU_BLENDFACTOR_ONE,
			//	SDL_GPU_BLENDFACTOR_ZERO,
			//	SDL_GPU_BLENDOP_ADD,
			//	0xf,
			//	true,
			//	false
			//};
		}
		// BLENDMODE_BLEND
		{
			auto& descs = statePrefabs.colorTargetDescriptionsPool.emplace_back(init);
			descs[0].blend_state = {
				SDL_GPU_BLENDFACTOR_SRC_ALPHA,
				SDL_GPU_BLENDFACTOR_ONE_MINUS_SRC_ALPHA,
				SDL_GPU_BLENDOP_ADD,
				SDL_GPU_BLENDFACTOR_ONE,
				SDL_GPU_BLENDFACTOR_ONE_MINUS_SRC_ALPHA,
				SDL_GPU_BLENDOP_ADD,
				0xf,
				true,
				false
			};
		}
		// BLENDMODE_BLEND_PREMULTIPLIED
		{
			auto& descs = statePrefabs.colorTargetDescriptionsPool.emplace_back(init);
			descs[0].blend_state = {
				SDL_GPU_BLENDFACTOR_ONE,
				SDL_GPU_BLENDFACTOR_ONE_MINUS_SRC_ALPHA,
				SDL_GPU_BLENDOP_ADD,
				SDL_GPU_BLENDFACTOR_ONE,
				SDL_GPU_BLENDFACTOR_ONE_MINUS_SRC_ALPHA,
				SDL_GPU_BLENDOP_ADD,
				0xf,
				true,
				false
			};
		}
		// BLENDMODE_ADD
		{
			auto& descs = statePrefabs.colorTargetDescriptionsPool.emplace_back(init);
			descs[0].blend_state = {
				SDL_GPU_BLENDFACTOR_SRC_ALPHA,
				SDL_GPU_BLENDFACTOR_ONE,
				SDL_GPU_BLENDOP_ADD,
				SDL_GPU_BLENDFACTOR_ZERO,
				SDL_GPU_BLENDFACTOR_ONE,
				SDL_GPU_BLENDOP_ADD,
				0xf,
				true,
				false
			};
		}
		// BLENDMODE_ADD_PREMULTIPLIED
		{
			auto& descs = statePrefabs.colorTargetDescriptionsPool.emplace_back(init);
			descs[0].blend_state = {
				SDL_GPU_BLENDFACTOR_ONE,
				SDL_GPU_BLENDFACTOR_ONE,
				SDL_GPU_BLENDOP_ADD,
				SDL_GPU_BLENDFACTOR_ZERO,
				SDL_GPU_BLENDFACTOR_ONE,
				SDL_GPU_BLENDOP_ADD,
				0xf,
				true,
				false
			};
		}
		// BLENDMODE_MOD
		{
			auto& descs = statePrefabs.colorTargetDescriptionsPool.emplace_back(init);
			descs[0].blend_state = {
				SDL_GPU_BLENDFACTOR_ZERO,
				SDL_GPU_BLENDFACTOR_SRC_COLOR,
				SDL_GPU_BLENDOP_ADD,
				SDL_GPU_BLENDFACTOR_ZERO,
				SDL_GPU_BLENDFACTOR_ONE,
				SDL_GPU_BLENDOP_ADD,
				0xf,
				true,
				false
			};
		}
		// BLENDMODE_MUL
		{
			auto& descs = statePrefabs.colorTargetDescriptionsPool.emplace_back(init);
			descs[0].blend_state = {
				SDL_GPU_BLENDFACTOR_DST_COLOR,
				SDL_GPU_BLENDFACTOR_ONE_MINUS_SRC_ALPHA,
				SDL_GPU_BLENDOP_ADD,
				SDL_GPU_BLENDFACTOR_ZERO,
				SDL_GPU_BLENDFACTOR_ONE,
				SDL_GPU_BLENDOP_ADD,
				0xf,
				true,
				false
			};
		}
	}

	bool GPUGraphicsPipelinePool::_CreateGraphicsPipelineState(GPLState state) {
		auto& vertexShader = statePrefabs.vertexShaders[state.vertexShaderID];
		auto& fragmentShader = statePrefabs.fragmentShaders[state.fragmentShaderID];
		auto& targetInfo = statePrefabs.targetInfos[state.targetInfoID];
		auto& colorTargetDescriptions = statePrefabs.colorTargetDescriptionsPool[state.colorTargetDescriptionsID];
		targetInfo.color_target_descriptions = colorTargetDescriptions.data();
		targetInfo.num_color_targets = (Uint32)colorTargetDescriptions.size();

		GPUGraphicsPipelineCreateInfo createInfo{
			*vertexShader,
			*fragmentShader,
			{
				vertexShader->vertexBufferDescriptions.data(),
				(Uint32)vertexShader->vertexBufferDescriptions.size(),
				vertexShader->vertexAttributes.data(),
				(Uint32)vertexShader->vertexAttributes.size()
			},
			(GPUPrimitiveType)state.primitiveType,
			statePrefabs.rasterizerStates[state.rasterizerStateID],
			statePrefabs.multisampleStates[state.multisampleStateID],
			statePrefabs.depthStencilStates[state.depthStencilStateID],
			targetInfo
		};

		if (!*pipelines.emplace_back(std::make_shared<GPUGraphicsPipeline>(device, createInfo))) {
			pipelines.pop_back();
			return false;
		}
		states.push_back(state);
		return true;
	}

	bool GPUGraphicsPipelinePool::_CreateGraphicsPipelines() {
		const GPLState pipelineStates[] = {
			{ SDL_GPU_PRIMITIVETYPE_POINTLIST,		VS_POINTVERTEX,		FS_OUTCOLOR			},
			{ SDL_GPU_PRIMITIVETYPE_POINTLIST,		VS_POINTVERTEX,		FS_CIRCLEPOINT		},
			{ SDL_GPU_PRIMITIVETYPE_POINTLIST,		VS_POINTVERTEX,		FS_AACIRCLEPOINT,	CTD_BLENDMODE_BLEND	},
			{ SDL_GPU_PRIMITIVETYPE_LINELIST,		VS_POSITIONVERTEX,	FS_OUTCOLOR			},
			{ SDL_GPU_PRIMITIVETYPE_LINESTRIP,		VS_POSITIONVERTEX,	FS_OUTCOLOR			},
			{ SDL_GPU_PRIMITIVETYPE_TRIANGLELIST,	VS_LINEVERTEX,		FS_LINE				},
			{ SDL_GPU_PRIMITIVETYPE_TRIANGLELIST,	VS_LINEVERTEX,		FS_AALINE,			CTD_BLENDMODE_BLEND	},
			{ SDL_GPU_PRIMITIVETYPE_TRIANGLELIST,	VS_CIRCLEVERTEX,	FS_CIRCLE			},
			{ SDL_GPU_PRIMITIVETYPE_TRIANGLELIST,	VS_CIRCLEVERTEX,	FS_AACIRCLE,		CTD_BLENDMODE_BLEND	},
			{ SDL_GPU_PRIMITIVETYPE_TRIANGLELIST,	VS_POSITIONVERTEX,	FS_OUTCOLOR			},
			{ SDL_GPU_PRIMITIVETYPE_TRIANGLESTRIP,	VS_POSITIONVERTEX,	FS_OUTCOLOR			},
			{ SDL_GPU_PRIMITIVETYPE_TRIANGLELIST,	VS_TEXTUREVERTEX,	FS_OUTTEXTURE		},
			{ SDL_GPU_PRIMITIVETYPE_TRIANGLELIST,	VS_TEXTUREVERTEX,	FS_DOTTEXTURECOLOR	}
		};
		for (auto& state : pipelineStates) {
			if (!_CreateGraphicsPipelineState(state)) {
				return false;
			}
		}
		return true;
	}

	GPUComputePipelinePool::GPUComputePipelinePool(GPUDevice& device) : device(device) {
		workers.reserve(CPL_NUM);
		pipelines.reserve(CPL_NUM);
		_CreateComputePipelines();
	}

	void GPUComputePipelinePool::_CreateComputePipelines() {
		static const char* CSFilename[CPL_NUM] = {
			"FilterGrayscale.comp",
			"FilterInvertColor.comp",
			"FilterInvertMono.comp",
			"DrawAlphaPattern.comp"
		};
		SDLFileStorage storage("shaders");
		if (!storage) {
			return;
		}
		string ext;
		stringv driver = device.GetDeviceDriver();
		stringv entrypoint = "main";
		GPUShaderFormat format = SDL_GPU_SHADERFORMAT_INVALID;
		if (driver == "vulkan") {
			format = SDL_GPU_SHADERFORMAT_SPIRV;
			ext = ".spv";
		}
		else if (driver == "direct3d12") {
			format = SDL_GPU_SHADERFORMAT_DXIL;
			ext = ".dxil";
		}
		else {
			entrypoint = "main0";
			format = SDL_GPU_SHADERFORMAT_METALLIB;
			ext = ".metallib";
		}

		GPUComputeShaderInfo shaderInfo;
		shaderInfo.entrypoint = entrypoint;
		shaderInfo.format = format;

		// CPL_FILTER_GRAYSCALE
		{
			shaderInfo.numReadwriteStorageTextures = 1;
			shaderInfo.numUniformBuffers = 1;		// push Rect2D
			shaderInfo.countX = 8;
			shaderInfo.countY = 8;
			if (!storage.ReadFile(CSFilename[CPL_FILTER_GRAYSCALE] + ext, shaderInfo.code)) {
				return;
			}
			if (!*pipelines.emplace_back(std::make_shared<GPUComputePipeline>(device, shaderInfo))) {
				pipelines.pop_back();
				return;
			}
			workers.emplace_back(8, 8);
		}
		// CPL_FILTER_INVERTCOLOR
		{
			shaderInfo.numReadwriteStorageTextures = 1;
			shaderInfo.numUniformBuffers = 1;		// push Rect2D
			shaderInfo.countX = 8;
			shaderInfo.countY = 8;
			if (!storage.ReadFile(CSFilename[CPL_FILTER_INVERTCOLOR] + ext, shaderInfo.code)) {
				return;
			}
			if (!*pipelines.emplace_back(std::make_shared<GPUComputePipeline>(device, shaderInfo))) {
				pipelines.pop_back();
				return;
			}
			workers.emplace_back(8, 8);
		}
		// CPL_FILTER_INVERTMONO
		{
			shaderInfo.numReadwriteStorageTextures = 1;
			shaderInfo.numUniformBuffers = 1;		// push Rect2D
			shaderInfo.countX = 8;
			shaderInfo.countY = 8;
			if (!storage.ReadFile(CSFilename[CPL_FILTER_INVERTMONO] + ext, shaderInfo.code)) {
				return;
			}
			if (!*pipelines.emplace_back(std::make_shared<GPUComputePipeline>(device, shaderInfo))) {
				pipelines.pop_back();
				return;
			}
			workers.emplace_back(8, 8);
		}
		// CPL_DRAW_ALPHAPATTERN
		{
			shaderInfo.numReadwriteStorageTextures = 1;
			shaderInfo.numUniformBuffers = 2;			// push PatternComputeUniform, Rect2D+Uint32
			shaderInfo.numReadonlyStorageBuffers = 1;	// bind Uint8 buffer
			shaderInfo.countX = 8;
			shaderInfo.countY = 8;
			if (!storage.ReadFile(CSFilename[CPL_DRAW_ALPHAPATTERN] + ext, shaderInfo.code)) {
				return;
			}
			if (!*pipelines.emplace_back(std::make_shared<GPUComputePipeline>(device, shaderInfo))) {
				pipelines.pop_back();
				return;
			}
			workers.emplace_back(8, 8);
		}
	}
}
